import json
import os
from enum import Enum
from typing import Optional

from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import PydanticOutputParser, StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langchain_core.runnables import RunnablePassthrough
from pydantic import Field, BaseModel

os.environ["OPENAI_API_KEY"] = os.environ["OPENAI_API_KEY_ZHIHU"]
os.environ["OPENAI_API_BASE"] = os.environ["OPENAI_API_BASE_ZHIHU"]

llm = ChatOpenAI(model="gpt-4o",temperature=0)

class SortField(str,Enum):
    price = 'price'
    data = 'data'

class SortDir(str,Enum):
    ascend = 'ascend'
    descend = 'descend'

class Semantics(BaseModel):
    name:Optional[str] = Field(description="套餐名称",default=None)
    price_lower: Optional[float] = Field(description="价格下限", default=None)
    price_upper: Optional[float] = Field(description="价格上限", default=None)
    data_lower: Optional[float] = Field(description="流量下限", default=None)
    data_upper: Optional[float] = Field(description="流量上限", default=None)
    sort_foeld: Optional[SortField] = Field(description="按价格或流量排序", default=None)
    sort_dir:Optional[SortDir] = Field(description="升序还是降序",default=None)

propmt_template = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template("将用户的输入解析为指定的JSON输出，输出格式如下，\n{format_instruction}\n不要输出未提及的字段"),
    HumanMessagePromptTemplate.from_template("{input}")
])

# OutputParser
parser = PydanticOutputParser(pydantic_object=Semantics)
prompt = propmt_template.partial(format_instruction = parser.get_format_instructions())

#同步
chain = ({"input":RunnablePassthrough()} | prompt | llm | parser)
response = chain.invoke("不超过100元的流量大的套餐有哪些")
print(json.dumps(response.dict(),indent=4,ensure_ascii=False))

#流方
chain = ({"input":RunnablePassthrough()} | prompt | llm | StrOutputParser())
response = chain.stream("不超过100元的流量大的套餐有哪些")
for chunk in response:
    print(chunk,end='')